#pragma once

__forceinline bool FOctreeChildNodeSubset::Contains(FOctreeChildNodeRef ChildRef) const
{
	// This subset contains the child if it has all the bits set that are set for the subset containing only the child node.
	const FOctreeChildNodeSubset ChildSubset(ChildRef);
	return (ChildBits & ChildSubset.ChildBits) == ChildSubset.ChildBits;
}

__forceinline FOctreeChildNodeSubset FOctreeNodeContext::GetIntersectingChildren(const BoxCenterAndExtent& QueryBounds) const
{
	FOctreeChildNodeSubset Result;
	
	// Load the query bounding box values as VectorRegisters.
	const VectorRegister QueryBoundsCenter = VectorLoadAligned(&QueryBounds.Center);
	const VectorRegister QueryBoundsExtent = VectorLoadAligned(&QueryBounds.Extent);
	const VectorRegister QueryBoundsMax = VectorAdd(QueryBoundsCenter,QueryBoundsExtent);
	const VectorRegister QueryBoundsMin = VectorSubtract(QueryBoundsCenter,QueryBoundsExtent);

	// Compute the bounds of the node's children.
	const VectorRegister BoundsCenter = VectorLoadAligned(&Bounds.Center);
	const VectorRegister BoundsExtent = VectorLoadAligned(&Bounds.Extent);
	const VectorRegister PositiveChildBoundsMin = VectorSubtract(
		VectorAdd(BoundsCenter,VectorLoadFloat1(&ChildCenterOffset)),
		VectorLoadFloat1(&ChildExtent));
	const VectorRegister NegativeChildBoundsMax = VectorAdd(
		VectorSubtract(BoundsCenter,VectorLoadFloat1(&ChildCenterOffset)),
		VectorLoadFloat1(&ChildExtent));

	// Intersect the query bounds with the node's children's bounds.
	Result.bPositiveX = VectoryAnyGreaterThan(VectorReplicate(QueryBoundsMax,0),VectorReplicate(PositiveChildBoundsMin,0)) != false;
	Result.bPositiveY = VectoryAnyGreaterThan(VectorReplicate(QueryBoundsMax,1),VectorReplicate(PositiveChildBoundsMin,1)) != false;
	Result.bPositiveZ = VectoryAnyGreaterThan(VectorReplicate(QueryBoundsMax,2),VectorReplicate(PositiveChildBoundsMin,2)) != false;
	Result.bNegativeX = VectoryAnyGreaterThan(VectorReplicate(QueryBoundsMin,0),VectorReplicate(NegativeChildBoundsMax,0)) == false;
	Result.bNegativeY = VectoryAnyGreaterThan(VectorReplicate(QueryBoundsMin,1),VectorReplicate(NegativeChildBoundsMax,1)) == false;
	Result.bNegativeZ = VectoryAnyGreaterThan(VectorReplicate(QueryBoundsMin,2),VectorReplicate(NegativeChildBoundsMax,2)) == false;
	return Result;
}

__forceinline FOctreeChildNodeRef FOctreeNodeContext::GetContainingChild(const BoxCenterAndExtent& QueryBounds) const
{
	FOctreeChildNodeRef Result;
	
	// Load the query bounding box values as VectorRegisters.
	const VectorRegister QueryBoundsCenter = VectorLoadAligned(&QueryBounds.Center);
	const VectorRegister QueryBoundsExtent = VectorLoadAligned(&QueryBounds.Extent);

	// Compute the bounds of the node's children.
	const VectorRegister BoundsCenter = VectorLoadAligned(&Bounds.Center);
	const VectorRegister ChildCenterOffsetVector = VectorLoadFloat1(&ChildCenterOffset);
	const VectorRegister NegativeCenterDifference = VectorSubtract(QueryBoundsCenter,VectorSubtract(BoundsCenter,ChildCenterOffsetVector));
	const VectorRegister PositiveCenterDifference = VectorSubtract(VectorAdd(BoundsCenter,ChildCenterOffsetVector),QueryBoundsCenter);

	// If the query bounds isn't entirely inside the bounding box of the child it's closest to, it's not contained by any of the child nodes.
	const VectorRegister MinDifference = VectorMin(PositiveCenterDifference,NegativeCenterDifference);
	if (VectoryAnyGreaterThan(VectorAdd(QueryBoundsExtent,MinDifference),VectorLoadFloat1(&ChildExtent)))
	{
		Result.bNULL = true;
	}
	else
	{
		// Return the child node that the query is closest to as the containing child.
		Result.X = VectoryAnyGreaterThan(VectorReplicate(QueryBoundsCenter,0),VectorReplicate(BoundsCenter,0)) != false;
		Result.Y = VectoryAnyGreaterThan(VectorReplicate(QueryBoundsCenter,1),VectorReplicate(BoundsCenter,1)) != false;
		Result.Z = VectoryAnyGreaterThan(VectorReplicate(QueryBoundsCenter,2),VectorReplicate(BoundsCenter,2)) != false;
	}

	return Result;
}

template<typename ElementType,typename OctreeSemantics>
void TOctree<ElementType,OctreeSemantics>::AddElement(ElementType& Element)
{
	AddElementToNode(Element,RootNode,RootNodeContext);
}
															
template<typename ElementType,typename OctreeSemantics>
void TOctree<ElementType,OctreeSemantics>::AddElementToNode(
	ElementType& Element,
	const FNode& InNode,
	const FOctreeNodeContext& InContext
	)
{
	const BoxCenterAndExtent ElementBounds = OctreeSemantics::GetBoundingBox(Element);

	for (TConstIterator NodeIt(InNode,InContext); NodeIt.HasPendingNodes(); NodeIt.Advance())
	{
		const FNode& Node = NodeIt.GetCurrentNode();
		const FOctreeNodeContext& Context = NodeIt.GetCurrentContext();

		bool bAddElementToThisNode = false;

		// Increment the number of elements included in this node and its children.
		Node.InclusiveNumElements++;

		if (Node.IsLeaf())
		{
			// If this is a leaf, check if adding this element would turn it into a node by overflowing its element list.
			if (Node.Elements.size() + 1 > OctreeSemantics::MaxElementsPerLeaf && Context.Bounds.Extent.X > MinLeafExtent)
			{
				// Copy the leaf's elements, remove them from the leaf, and turn it into a node.
				ElementArrayType ChildElements;
				std::swap(ChildElements,Node.Elements);
				Node.InclusiveNumElements = 0;

				// Allow elements to be added to children of this node.
				Node.bIsLeaf = false;

				// Re-add all of the node's child elements, potentially creating children of this node for them.
				for (ElementArrayType::iterator it(ChildElements.begin()), end(ChildElements.end()); it != end; ++it)
					AddElementToNode(*it,Node,Context);

				// Add the element to this node.
				AddElementToNode(Element,Node,Context);
				return;
			}
			else
			{
				// If the leaf has room for the new element, simply add it to the list.
				bAddElementToThisNode = true;
			}
		}
		else
		{
			// If this isn't a leaf, find a child that entirely contains the element.
			const FOctreeChildNodeRef ChildRef = Context.GetContainingChild(ElementBounds);	
			if (ChildRef.IsNULL())
			{
				// If none of the children completely contain the element, add it to this node directly.
				bAddElementToThisNode = true;
			}
			else
			{
				// Create the child node if it hasn't been created yet.
				if(!Node.Children[ChildRef.Index])
				{
					Node.Children[ChildRef.Index] = new typename TOctree<ElementType,OctreeSemantics>::FNode(&Node);
				}

				// Push the node onto the stack to visit.
				NodeIt.PushChild(ChildRef);
			}
		}

		if (bAddElementToThisNode)
		{
			if (Node.Elements.capacity() == 0)
				Node.Elements.reserve(OctreeSemantics::MaxElementsPerLeaf);

			// Add the element to this node.
			Node.Elements.push_back(Element);
			
			// Set the element's ID.
			OctreeSemantics::SetElementId(Element,FOctreeElementId(&Node, static_cast<int>(Node.Elements.size() - 1)));
			return;
		}
	}
}

template<typename ElementType,typename OctreeSemantics>
void TOctree<ElementType,OctreeSemantics>::RemoveElement(FOctreeElementId ElementId)
{
	assert(ElementId.IsValidId()); 

	FNode* ElementIdNode = (FNode*)ElementId.Node;

	// swap the element if it isn't at the end.
	if (ElementId.ElementIndex != ElementIdNode->Elements.size() - 1)
	{
		ElementIdNode->Elements[ElementId.ElementIndex] = ElementIdNode->Elements[ElementIdNode->Elements.size() - 1];

		// Update the external element id for the element that was swapped into the vacated element index.
		OctreeSemantics::SetElementId(ElementIdNode->Elements[ElementId.ElementIndex],ElementId);
	}
	ElementIdNode->Elements.resize(ElementIdNode->Elements.size() - 1);

	// Update the inclusive element counts for the nodes between the element and the root node,
	// and find the largest node that is small enough to collapse.
	const FNode* CollapseNode = NULL;
	for (const FNode* Node = ElementIdNode;Node;Node = Node->Parent)
	{
		--Node->InclusiveNumElements;
		if (Node->InclusiveNumElements < OctreeSemantics::MinInclusiveElementsPerNode)
		{
			CollapseNode = Node;
		}
	}

	// Collapse the largest node that was pushed below the threshold for collapse by the removal.
	if (CollapseNode == NULL)
		return;

	// Gather the elements contained in this node and its children.
	ElementArrayType CollapsedChildElements;
	CollapsedChildElements.reserve(CollapseNode->InclusiveNumElements);
	for (TConstIterator ChildNodeIt(*CollapseNode,RootNodeContext);ChildNodeIt.HasPendingNodes();ChildNodeIt.Advance())
	{
		const FNode& ChildNode = ChildNodeIt.GetCurrentNode();

		// Add the child's elements to the collapsed element list.
		for(ElementArrayType::iterator it(ChildNode.Elements.begin()), end(ChildNode.Elements.end()); it != end; ++it)
		{
			CollapsedChildElements.push_back(*it);

			// Update the external element id for the element that's being collapsed.
			OctreeSemantics::SetElementId(*it, FOctreeElementId(CollapseNode, CollapsedChildElements.size() - 1));
		}

		// Recursively visit all child nodes.
		FOREACH_OCTREE_CHILD_NODE(ChildRef)
		{
			if(ChildNode.HasChild(ChildRef))
			{
				ChildNodeIt.PushChild(ChildRef);
			}
		}
	}

	// Replace the node's elements with the collapsed element list.
	std::swap(CollapseNode->Elements,CollapsedChildElements);

	// Mark the node as a leaf.
	CollapseNode->bIsLeaf = true;

	// Free the child nodes.
	FOREACH_OCTREE_CHILD_NODE(ChildRef)
	{
		delete CollapseNode->Children[ChildRef.Index];
		CollapseNode->Children[ChildRef.Index] = NULL;
	}
}

template<typename ElementType,typename OctreeSemantics>
ElementType& TOctree<ElementType,OctreeSemantics>::GetElementById(FOctreeElementId ElementId)
{
	check(ElementId.IsValidId());
	FNode* ElementIdNode = (FNode*)ElementId.Node;
	return ElementIdNode->Elements(ElementId.ElementIndex);
}

template<typename ElementType,typename OctreeSemantics>
void TOctree<ElementType,OctreeSemantics>::DumpStats()
{
	INT NumNodes = 0;
	INT NumLeaves = 0;
	INT NumElements = 0;
	INT MaxElementsPerNode = 0;
	for (TConstIterator<> NodeIt(*this);NodeIt.HasPendingNodes();NodeIt.Advance())
	{
		const FNode& CurrentNode = NodeIt.GetCurrentNode();

		NumNodes++;
		if(CurrentNode.IsLeaf())
		{
			NumLeaves++;
		}

		NumElements += CurrentNode.GetElements().Num();
		MaxElementsPerNode = Max(MaxElementsPerNode,CurrentNode.GetElements().Num());

		FOREACH_OCTREE_CHILD_NODE(ChildRef)
		{
			if(CurrentNode.HasChild(ChildRef))
			{
				NodeIt.PushChild(ChildRef);
			}
		}
	}
}

template<typename ElementType,typename OctreeSemantics>
TOctree<ElementType,OctreeSemantics>::TOctree(const Vector3& InOrigin, float InExtent)
:	RootNode(NULL)
,	RootNodeContext(BoxCenterAndExtent(InOrigin,Vector3(InExtent,InExtent,InExtent)))
,	MinLeafExtent(InExtent * pow((1.0f + 1.0f / (float)FOctreeNodeContext::LoosenessDenominator) / 2.0f,OctreeSemantics::MaxNodeDepth))
{
}
